from tensorflow.keras.layers import Input, Activation, Conv2D, ReLU, Flatten, Dense, Reshape, Conv2DTranspose, BatchNormalization, Dropout
import math

# 定义编码器
def build_encoder(input_tensor, feature_size):
    x = Conv2D(32, (3, 3), strides=1, padding='same')(input_tensor)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = ReLU()(x)
    x = Conv2D(64, (3, 3), strides=1, padding='same')(x)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = ReLU()(x)
    x = Conv2D(128, (3, 3), strides=1, padding='same')(x)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = ReLU()(x)
    encoded = Flatten()(x)
    encoded = Dropout(0.2)(encoded)
    encoded = Dense(feature_size, name='encoder_output_layer')(encoded)
    return encoded

# 定义解码器
def build_decoder(encoded):
    minsize = 2
    x = Dense(512)(encoded)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = Reshape((minsize, minsize, math.floor(512 / (minsize * minsize))))(x)
    x = Conv2DTranspose(128, (3, 3), strides=2, padding='same')(x)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = ReLU()(x)
    x = Conv2DTranspose(64, (3, 3), strides=2, padding='same')(x)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = ReLU()(x)
    x = Conv2DTranspose(32, (3, 3), strides=2, padding='same')(x)
    x = BatchNormalization()(x)  # 添加批量归一化
    x = ReLU()(x)
    decoded = Conv2DTranspose(1, (3, 3), strides=1, padding='same', output_padding=0)(x)
    decoded = BatchNormalization()(decoded)  # 添加批量归一化
    decoded = Activation('tanh', name='decoder_output_layer')(decoded)
    return decoded

# 定义分类器
def build_classifier(encoded, n):
    classifier_output = Dense(units=n, activation='softmax', name='classifier_output_layer')(encoded)
    return classifier_output